#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <string.h>
#include <math.h>

#include "./svpng/svpng.inc"

class Maze
{
private:
    // 方块
    struct Block
    {
        Block *next = this;  // “下一个”方块，用于识别不同连通片
        bool right = false;  // 右边是否有相邻方块
        bool down = false;   // 下方是否有相邻方块
        bool marked = false; // 是否被标记
    };

    // 边
    struct Edge
    {
        enum Type
        {
            H, // 水平
            V  // 垂直
        } type;
        int rnd = rand();       // 随机编号
        Block *side1 = nullptr; // 指向上方（水平）或左边（垂直）方块
        Block *side2 = nullptr; // 指向下方（水平）或右边（垂直）方块
    };

    int width, height; // 宽，高
    Block **map;       // 地图

    // 查找连通片
    Block *find(Block *block)
    {
        if (block != block->next)
        {
            block->next = find(block->next);
        }
        return block->next;
    }

    // 获取相邻方块数量
    int neighbors(int x, int y)
    {
        int a = map[y][x].right && map[y][x + 1].marked;
        int b = map[y][x].down && map[y + 1][x].marked;
        int c = x && map[y][x - 1].right && map[y][x - 1].marked;
        int d = y && map[y - 1][x].down && map[y - 1][x].marked;
        return a + b + c + d;
    }

    // 判断一个方块是否为通路
    bool is_pass(int x, int y)
    {
        // int a = map[y][x].right && map[y][x + 1].marked;
        // int b = map[y][x].down && map[y + 1][x].marked;
        // int c = x && map[y][x - 1].right && map[y][x - 1].marked;
        // int d = y && map[y - 1][x].down && map[y - 1][x].marked;
        // return a + b + c + d == 1;
        return neighbors(x, y) == 1;
    }

    // 判断一个方块是否为道路终点
    bool is_end(int x, int y)
    {
        return map[y][x].right + map[y][x].down + (x && map[y][x - 1].right) + (y && map[y - 1][x].down) <= 1;
    }

    // 下一条通路
    void next_pass(int &x, int &y)
    {
        if (map[y][x].right && map[y][x + 1].marked)
        {
            x++;
        }
        else if (map[y][x].down && map[y + 1][x].marked)
        {
            y++;
        }
        else if (x && map[y][x - 1].right && map[y][x - 1].marked)
        {
            x--;
        }
        else if (y && map[y - 1][x].down && map[y - 1][x].marked)
        {
            y--;
        }
        else
        {
            return;
        }
    }

public:
    // 初始化
    Maze(int width, int height)
    {
        this->width = width;
        this->height = height;

        this->map = new Block *[height];

        for (int i = 0; i < height; i++)
        {
            map[i] = new Block[width];
        }
    }

    // 删除地图，释放内存空间
    ~Maze()
    {
        for (int i = 0; i < height; i++)
        {
            delete map[i];
        }
        delete map;
    }

    int get_width() { return width; }

    int get_height() { return height; }

    // 重置
    void reset()
    {
        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                map[i][j].right = false;
                map[i][j].down = false;
                map[i][j].next = &map[i][j];
                map[i][j].marked = false;
            }
        }
    }

    void generate() // 生成地图
    {
        int num_edges = 2 * width * height - width - height;
        Edge **edges = new Edge *[num_edges];

        int idx = 0;

        // 添加水平边
        for (int i = 0; i < height - 1; i++)
        {
            for (int j = 0; j < width; j++)
            {
                Edge *hor = new Edge;
                hor->type = Edge::Type::H;
                hor->side1 = &map[i][j];
                hor->side2 = &map[i + 1][j];
                edges[idx++] = hor;
            }
        }

        // 添加垂直边
        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width - 1; j++)
            {
                Edge *ver = new Edge;
                ver->type = Edge::Type::V;
                ver->side1 = &map[i][j];
                ver->side2 = &map[i][j + 1];
                edges[idx++] = ver;
            }
        }

        // 打乱边
        qsort(edges, num_edges, sizeof(Edge *), [](const void *e1, const void *e2)
              { return (*(Edge **)e1)->rnd - (*(Edge **)e2)->rnd; });

        // 随机删除边
        for (int i = 0; i < num_edges; i++)
        {
            Edge *edge = edges[i];
            if (find(edge->side1) != find(edge->side2))
            {
                if (edge->type == Edge::Type::H)
                {
                    edge->side1->down = true;
                }
                else
                {
                    edge->side1->right = true;
                }
                edge->side1->next->next = edge->side2->next;
            }
            delete edge;
        }

        delete edges;
    }

    // 寻找路径
    void path(int x1, int y1, int x2, int y2)
    {
        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                map[i][j].marked = true; // 先全部标记
            }
        }

        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                if (!(i == y1 && j == x1 || i == y2 && j == x2) && map[i][j].marked && is_end(j, i))
                {
                    int end_x = j, end_y = i;

                    while (!(end_y == y1 && end_x == x1 || end_y == y2 && end_x == x2) && is_pass(end_x, end_y))
                    {
                        map[end_y][end_x].marked = false;
                        next_pass(end_x, end_y);
                    }
                }
            }
        }
    }

    void show()
    {
        printf("  ");
        for (int j = 0; j < width; j++)
        {
            printf("%c%c", '_', j < width - 1 ? '_' : ' ');
        }
        printf("\n");
        for (int i = 0; i < height; i++)
        {
            printf(" |");
            for (int j = 0; j < width; j++)
            {
                printf("%c%c", map[i][j].marked ? '*' : (map[i][j].down ? ' ' : '_'), map[i][j].right ? (map[i][j].down ? ' ' : '_') : '|');
            }
            printf("\n");
        }
    }

    void save_bin(const char *filename)
    {
        FILE *fp = fopen(filename, "wb");
        if (!fp)
        {
            return;
        }

        fwrite(&width, sizeof(width), 1, fp);
        fwrite(&height, sizeof(height), 1, fp);
        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                char data = map[i][j].right << 4 | map[i][j].down;
                fwrite(&data, sizeof(data), 1, fp);
            }
        }

        fclose(fp);
    }

    static Maze *open_bin(const char *filename)
    {
        FILE *fp = fopen(filename, "rb");
        if (!fp)
        {
            return nullptr;
        }

        int width, height;
        fread(&width, sizeof(width), 1, fp);
        fread(&height, sizeof(height), 1, fp);

        Maze *maze = new Maze(width, height);
        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                char data;
                fread(&data, sizeof(data), 1, fp);
                maze->map[i][j].right = data >> 4;
                maze->map[i][j].down = data & 0xf;
            }
        }

        fclose(fp);
        return maze;
    }

    void save_img(const char *filename)
    {
        struct RGB
        {
            unsigned char r;
            unsigned char g;
            unsigned char b;
        };

        RGB black = {0x0, 0x0, 0x0};
        RGB white = {0xff, 0xff, 0xff};
        RGB highlight = {0xff, 0x0, 0x0};

        FILE *fp = fopen(filename, "wb");
        if (!fp)
        {
            return;
        }

        int img_width = width * 2 + 1;
        int img_height = height * 2 + 1;

        RGB *rgb = new RGB[img_width * img_height];

        for (int j = 0; j < img_width; j++)
        {
            rgb[j] = black;
        }

        for (int i = 1; i < img_height; i++)
        {
            rgb[i * img_width] = black;
        }

        for (int i = 0; i < height; i++)
        {
            for (int j = 0; j < width; j++)
            {
                RGB color = map[i][j].marked ? highlight : white;
                rgb[(i * 2 + 1) * img_width + (j * 2 + 1)] = color;
                rgb[(i * 2 + 1) * img_width + (j * 2 + 2)] = map[i][j].right ? (map[i][j + 1].marked ? color : white) : black;
                rgb[(i * 2 + 2) * img_width + (j * 2 + 1)] = map[i][j].down ? (map[i + 1][j].marked ? color : white) : black;
                rgb[(i * 2 + 2) * img_width + (j * 2 + 2)] = black;
            }
        }

        svpng(fp, img_width, img_height, (unsigned char *)rgb, 0);
        delete rgb;
        fclose(fp);
    }
};

void help(const char *exe_name)
{
    const char *help_message =
        "Usage: %s [options] [<filename>]\n"
        "Options:\n"
        "  -h --help           show this help message\n"
        "  -s --show           show maze in ASCII characters\n"
        "  -g --generate size  generate a maze in specified size e.g. 20x30\n"
        "  -b --save-binary    save maze to binary file\n"
        "  -i --save-image     save maze to png file\n"
        "  -p --path           find path between two locations e.g. (0,0)-(10,20)\n";
    printf(help_message, exe_name);
}

int main(int argc, const char *argv[])
{
    const char *exe_name;
    for (exe_name = argv[0] + strlen(argv[0]); exe_name > argv[0]; exe_name--)
    {
        if (*exe_name == '/' || *exe_name == '\\')
        {
            exe_name++;
            break;
        }
    }

    const char *filename = nullptr;
    const char *maze_size = nullptr;
    const char *save_binary = nullptr;
    const char *save_image = nullptr;
    const char *path = nullptr;
    bool show_maze = false;

    for (int i = 1; i < argc; i++)
    {
        const char *arg = argv[i];
        if (!strcmp(arg, "-h") || !strcmp(arg, "--help"))
        {
            help(exe_name);
            return 0;
        }
        else if (!strcmp(arg, "-s") || !strcmp(arg, "--show"))
        {
            show_maze = true;
        }
        else if (!strcmp(arg, "-g") || !strcmp(arg, "--generate"))
        {
            maze_size = argv[++i];
        }
        else if (!strcmp(arg, "-b") || !strcmp(arg, "--save-binary"))
        {
            save_binary = argv[++i];
        }
        else if (!strcmp(arg, "-i") || !strcmp(arg, "--save-image"))
        {
            save_image = argv[++i];
        }
        else if (!strcmp(arg, "-p") || !strcmp(arg, "--path"))
        {
            path = argv[++i];
        }
        else
        {
            filename = arg;
        }
    }

    Maze *maze;
    if (filename)
    {
        maze = Maze::open_bin(filename);
    }
    else if (maze_size)
    {
        int width = -1, height = -1;
        sscanf(maze_size, "%dx%d", &width, &height);

        if (!(width > 0 && height > 0))
        {
            exit(0);
        }
        maze = new Maze(width, height);

        srand(time(0));
        maze->generate();
    }
    else
    {
        help(exe_name);
        return 0;
    }

    if (path)
    {
        int x1 = -1, y1 = -1, x2 = -1, y2 = -1;
        sscanf(path, "(%d,%d)-(%d,%d)", &x1, &y1, &x2, &y2);

        if (!(x1 >= 0 && x1 < maze->get_width() && y1 >= 0 && y1 < maze->get_height() && x2 >= 0 && x2 < maze->get_width() && y2 >= 0 && y2 < maze->get_height()))
        {
            exit(0);
        }

        maze->path(x1, y1, x2, y2);
    }
    if (show_maze)
    {
        maze->show();
    }
    if (save_binary)
    {
        maze->save_bin(save_binary);
    }
    if (save_image)
    {
        maze->save_img(save_image);
    }

    delete maze;
    return 0;
}
